from flask import Flask, request, render_template, jsonify, send_from_directory
from werkzeug.utils import secure_filename
import os
import json
import redis
import time
from datetime import datetime

app = Flask(__name__)

# Redis 配置
redis_client = redis.Redis(host='localhost', port=6379, db=0)

# 文件保存路径
UPLOAD_FOLDER = os.path.join(app.root_path, 'uploads')
OUTPUT_FOLDER = os.path.join(app.root_path, 'outputs')
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/uploads/<filename>')
def uploaded_file(filename):
    """提供上传文件的访问"""
    return send_from_directory(UPLOAD_FOLDER, filename)

@app.route('/outputs/<filename>')
def output_file(filename):
    """提供输出文件的访问"""
    return send_from_directory(OUTPUT_FOLDER, filename)

@app.route('/upload', methods=['POST'])
def upload():
    if 'image' not in request.files:
        return jsonify({'error': 'No image uploaded'}), 400
    
    image = request.files['image']
    if image.filename == '':
        return jsonify({'error': 'No image selected'}), 400
    
    # 获取选择的工作流
    workflow_id = request.form.get('workflow_id', 'default')
    workflow_path = os.path.join(app.root_path, 'workflow', f'{workflow_id}.json')
    
    if not os.path.exists(workflow_path):
        return jsonify({'error': 'Workflow not found'}), 404
    
    # 保存上传的图片
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    task_id = f'task_{timestamp}'
    image_filename = f'input_{timestamp}.png'
    image_path = os.path.join(UPLOAD_FOLDER, image_filename)
    image.save(image_path)
    
    # 读取工作流配置
    with open(workflow_path, 'r') as f:
        workflow_data = json.load(f)
    
    # 创建任务数据
    task_data = {
        'task_id': task_id,
        'status': 'pending',
        'image_path': image_filename,
        'workflow': workflow_data,
        'workflow_id': workflow_id,
        'created_at': time.time()
    }
    
    # 将任务添加到Redis队列
    redis_client.hset('tasks', task_id, json.dumps(task_data))
    redis_client.rpush('task_queue', task_id)
    
    return jsonify({
        'task_id': task_id,
        'status': 'pending',
        'workflow_id': workflow_id
    })

@app.route('/workflows')
def list_workflows():
    """列出所有可用的工作流"""
    workflow_dir = os.path.join(app.root_path, 'workflow')
    workflows = []
    for filename in os.listdir(workflow_dir):
        if filename.endswith('.json'):
            workflow_path = os.path.join(workflow_dir, filename)
            with open(workflow_path, 'r') as f:
                try:
                    workflow_data = json.load(f)
                    workflows.append({
                        'id': filename[:-5],  # 移除 .json 后缀
                        'name': filename,
                        'path': workflow_path
                    })
                except json.JSONDecodeError:
                    continue
    return jsonify(workflows)

@app.route('/status/<task_id>')
def get_status(task_id):
    task_data = redis_client.hget('tasks', task_id)
    if not task_data:
        return jsonify({'error': 'Task not found'}), 404
    
    task_info = json.loads(task_data)
    return jsonify(task_info)

@app.route('/task/pull', methods=['POST'])
def pull_task():
    """供 ServerA 拉取任务的接口"""
    task_id = redis_client.lpop('task_queue')
    if not task_id:
        return jsonify({'status': 'no_task'}), 404
    
    task_id = task_id.decode('utf-8')
    task_data = redis_client.hget('tasks', task_id)
    if not task_data:
        return jsonify({'error': 'Task data not found'}), 404
    
    task_info = json.loads(task_data)
    # 更新任务状态
    task_info['status'] = 'processing'
    redis_client.hset('tasks', task_id, json.dumps(task_info))
    
    return jsonify({
        'task_id': task_id,
        'image_path': os.path.join(UPLOAD_FOLDER, task_info['image_path']),
        'workflow': task_info['workflow']
    })

@app.route('/task/complete', methods=['POST'])
def complete_task():
    """供 ServerA 回传处理结果的接口"""
    data = request.json
    task_id = data.get('task_id')
    if not task_id:
        return jsonify({'error': 'Task ID required'}), 400
    
    task_data = redis_client.hget('tasks', task_id)
    if not task_data:
        return jsonify({'error': 'Task not found'}), 404
    
    task_info = json.loads(task_data)
    task_info['status'] = 'completed'
    task_info['result'] = data.get('result')
    task_info['completed_at'] = time.time()
    
    redis_client.hset('tasks', task_id, json.dumps(task_info))
    return jsonify({'status': 'success'})

@app.route('/save_output', methods=['POST'])
def save_output():
    """保存 ComfyUI 生成的输出图片"""
    if 'image' not in request.files:
        return jsonify({'error': 'No image file'}), 400
        
    image_file = request.files['image']
    if image_file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
        
    # 保存图片
    filename = secure_filename(image_file.filename)
    filepath = os.path.join(OUTPUT_FOLDER, filename)
    image_file.save(filepath)
    
    return jsonify({'filename': filename}), 200

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=15000)
